#include "mmu_config.h"
#include "mmu.h"
#include "svma.h"

#define TALBE_SIZE 0x1000   // the table size when the granule size is 4k

extern mx_void_pt __stack_base;
extern mx_void_pt __start;

/* the implementation of `mmu_table_alloc_tg4k` */
mx_void_pt mmu_table_alloc_tg4k ()
{
    return mx_svma_heap_alloc();
}

static mx_uint64_t ttbr0;
static mx_uint64_t ttbr1;

/* implementation of `mx_svma_mmuconfig` */
mx_usize_t mx_svma_mmuconfig ()
{
    /* TTBR initialization */
    ttbr0 = (mx_uint64_t)mx_svma_tbm_alloc(TALBE_SIZE);
    ttbr1 = (mx_uint64_t)mx_svma_tbm_alloc(TALBE_SIZE);
    asm volatile("msr ttbr0_el1, %0\n"::"r"(ttbr0));
    asm volatile("msr ttbr1_el1, %0\n dsb sy\n"::"r"(ttbr1));
    /* MAIR configuration */
    mx_uint64_t regv = MMU_MAIR_NORMAL(0, _NORMAL_WB_TRANSIENT(1, 1), _NORMAL_WB_NON_TRANSIENT(1, 1)) |
                       MMU_MAIR_NORMAL(1, _NORMAL_NON_CACHEABLE, _NORMAL_NON_CACHEABLE);
    asm volatile("msr mair_el1, %0\n dsb sy\n"::"r"(regv));
    /* tcr configuration, granule size: 4kB, 48-bit virtual address */
    regv = 0;
    FILMSKS(regv, MMU_TCR_T0SZ_MASK, 16);       /* TTBR0, T0SZ = 16, the bit numbers of virtual address is 48 (64 - T0SZ) */
    CLSMSK(regv, MMU_TCR_EPD0_MASK);            /* TTBR0, EPD0 = 0 */
    FILMSKS(regv, MMU_TCR_IRGN0_MASK, 3);       /* TTBR0, memory attribute: Normal memory, Inner Write-Back no Write-Allocate Cacheable */
    FILMSKS(regv, MMU_TCR_ORGN0_MASK, 3);       /* TTBR0, memory attribute: Normal memory, Outer Write-Back no Write-Allocate Cacheable */
    FILMSKS(regv, MMU_TCR_SH0_MASK, 2);         /* TTBR0, memory shareability: Outer shareable */
    FILMSKS(regv, MMU_TCR_TG0_MASK, 0);         /* TTBR0, granule size: 4kB, L0-bit[47:39] L1-bit[38:30] L2-bit[29-21] L3-bit[20:12] */
    FILMSKS(regv, MMU_TCR_T1SZ_MASK, 16);       /* TTBR1, T1SZ = 16, the bit numbers of virtual address is 48 (64 - T1SZ) */
    CLSMSK(regv, MMU_TCR_A1_MASK);              /* A1 = 0, use TTBR0.ASID, for user */
    CLSMSK(regv, MMU_TCR_EPD1_MASK);            /* TTBR1, EPD1 = 0 */
    FILMSKS(regv, MMU_TCR_IRGN1_MASK, 3);       /* TTBR1, memory attribute: Normal memory, Inner Write-Back no Write-Allocate Cacheable */
    FILMSKS(regv, MMU_TCR_ORGN1_MASK, 3);       /* TTBR1, memory attribute: Normal memory, Outer Write-Back no Write-Allocate Cacheable */
    FILMSKS(regv, MMU_TCR_SH1_MASK, 2);         /* TTBR1, memory shareability: Outer shareable */
    FILMSKS(regv, MMU_TCR_TG1_MASK, 2);         /* TTBR1, granule size: 4kB, L0-bit[47:39] L1-bit[38:30] L2-bit[29-21] L3-bit[20:12] */
    FILMSKS(regv, MMU_TCR_IPS_MASK, 1);         /* IPS = 1, 36-bit(64GB) */
    SETMSK(regv, MMU_TCR_AS_MASK);              /* AS = 1, ASID size is 16-bit */
    CLSMSK(regv, MMU_TCR_TBI0_MASK);            /* TTBR0, disable tag function */
    CLSMSK(regv, MMU_TCR_TBI1_MASK);            /* TTBR1, disable tag function */
    asm volatile("msr tcr_el1, %0\n"::"r"(regv));
    return (1 << 12);
}

/* implementation of `mx_svma_sysmap` */
void mx_svma_sysmap ()
{
    /* kernel map */
    mx_usize_t system_space = (mx_usize_t)&__stack_base - (mx_usize_t)&__start + 0x80000;
    system_space += MMU_ALIGN_4K_MASK;
    system_space &= ~MMU_ALIGN_4K_MASK;
    mmu_map_4k((mx_uint64_pt)ttbr0, (mx_void_pt)0x40080000, (mx_void_pt)0x40080000, mx_false, 1);
    mmu_map_4k((mx_uint64_pt)ttbr1, (mx_void_pt)0xffff000000000000, (mx_void_pt)0x40000000, mx_false, system_space / (1 << MMU_ALIGN_4K));
    /* GIC map */
    mmu_map_4k((mx_uint64_pt)ttbr0, (mx_void_pt)0x8000000U, (mx_void_pt)0x8000000U, mx_true, 0x10000 / (1 << MMU_ALIGN_4K));
    mmu_map_4k((mx_uint64_pt)ttbr0, (mx_void_pt)0x80a0000U, (mx_void_pt)0x80a0000U, mx_true, 0x20000 / (1 << MMU_ALIGN_4K));
    /* uart map */
    mmu_map_4k((mx_uint64_pt)ttbr0, (mx_void_pt)0x09000000, (mx_void_pt)0x09000000, mx_true, 0x1000 / (1 << MMU_ALIGN_4K));
}
